from functools import partial
import time
from collections import defaultdict

from osm_utils import *
from problem_generator import *
from map_problem import *
from search_runners import *
from aima.search import *


FORWARD = 0
BACKWARD = 1

def createCoolGraph(road_map, B):
    links = {}
    for j in road_map.junctions:
        for direction in (FORWARD, BACKWARD):
            links[(direction, j.key)] = [Link((direction, l.source), (direction, l.target), l.distance, l.highway_type, l.speed) 
                                         for l in j.links]
    links[(FORWARD, B)].append(Link((FORWARD, B), (BACKWARD, B), 0, 0, 1))
    return links


def createDijekstraGraph(road_map, B):
    links = defaultdict(list)
    for j in road_map.junctions:
        for direction in (FORWARD, BACKWARD):
            for l in j.links:
                links[(direction, l.target)].append(Link((direction, l.target), (direction, l.source), l.distance, l.highway_type, l.speed))
    links[(BACKWARD, B)].append(Link((BACKWARD, B), (FORWARD, B), 0, 0, 1))
    return links


from heapq import heappush, heappop
def dijkstra(N, S, edges, lim):
    '''https://gist.github.com/irachex/3922694/raw/9bd40836eee25215b30f83977a8a1831bbb3550a/dijkstra.py'''
    d = defaultdict(partial(float, "inf"))
    d[S] = 0
    h = []
    heappush(h, (0, S))
    for i in range(N - 1):
        min_dist, k = 0, 0
        if not h: break
        while h:
            min_dist, k = heappop(h)
            if min_dist == d[k]: break
        for l in edges[k]:
            w = l.CalcTime()
            j = l.target
            if d[j] > d[k] + w and lim >= d[k] + w:
                d[j] = d[k] + w
                heappush(h, (d[j], j))
    return d


def C_cool(problem, node1, link, next_key):
    for n in node1.path():
        prev_link = n.action
        if prev_link is not None and prev_link.source[1] == link.target[1] and prev_link.target[1] == link.source[1]:
            return prev_link.CalcTime() + link.CalcTime()
    return 0

def uniform_graph_search(problem, fringe):
    " based on aima.search.best_first_graphs_search "
    visited = {}
    fringe.append(Node(problem.initial))
    while fringe:
        node = fringe.pop()
        visited[node] = True
        if problem.goal_test(node): 
            return node, len(fringe) + len(visited)
        fringe.extend(node.expand(problem))
    return None, len(fringe) + len(visited)

class CoolMapProblem(Problem):
    """The abstract class for a formal problem.  You should subclass this and
    implement the method successor, and possibly __init__, goal_test, and
    path_cost. Then you will create instances of your subclass and solve them
    with the various search functions."""
    
    def __init__(self, initial_junction_key, goal_junction_key, links, dj_distances, limit, cost_func):
        self.initial = initial_junction_key
        self.goal = goal_junction_key
        self.links = links
        self.dj_distances = dj_distances
        self.limit = limit
        self.cost_func = partial(cost_func, self) # if cost_func else lambda node1, link, junction2: 1
        
    def successor(self, node):
        prev_time = sum(n.action.CalcTime() for n in node.path() if n.action is not None)
        return ((link, link.target) for link in self.links[node.state]
                if prev_time + self.dj_distances[link.target] <= self.limit and (link.target not in (n.state for n in node.path())))
    
    def goal_test(self, node):
        return (node.state == self.goal)    

    def path_cost(self, c, node1, link, junction2_key):
        """ Get's node1 and not junction1 because we sometimes need the path leading
        to junction1.. i.e. the C_s() func.."""
        return c + self.cost_func(node1, link, junction2_key)

    
def RunCoolPathSearch(rm, A, B):
    start_time = time.clock()

    node1, nodes_num1, run_time1, route_cost1 = RunAstar(rm, A, B, C_f, H_f)
    assert node1 is not None
    node2, nodes_num2, run_time2, route_cost2 = RunAstar(rm, B, A, C_f, H_f)
    assert node2 is not None

    Cf = route_cost1 + route_cost2;
    limit = 1.1 * Cf

    djlinks = createDijekstraGraph(rm, B)
    dj_distances = dijkstra(len(rm.junctions) * 2, (BACKWARD, A), djlinks, limit)
    assert(dj_distances[(FORWARD, A)] < limit)
    print 'dj_distances[(FORWARD, A)] = ', dj_distances[(FORWARD, A)]
    cool_links = createCoolGraph(rm, B)

    print 'RunCoolPathSearch {} -> {}, result: '.format(A, B)
    problem = CoolMapProblem((FORWARD, A), (BACKWARD, A), cool_links, dj_distances, limit, C_cool)
    
    def f(n):
        return n.path_cost
    #node, developed_nodes_count = search.best_first_graph_search(problem, f)
    node, developed_nodes_count = uniform_graph_search(problem, PriorityQueue(min, f))

    run_time = time.clock() - start_time
    # TRM debug?
    if node:
        path = node.path()
        print "\tdistance = {}, time = {}, stats = {}".format(
                                                sum(node.action.distance for node in path if node.action),
                                                sum(node.action.CalcTime() for node in path if node.action),
                                                problem)
    else:
        print "RunCoolPathSearch failed :(... (are they reachable??)"
    return node, developed_nodes_count, run_time, node.path_cost, Cf


if __name__ == "__main__":
    rm = CountryMap().LoadMap()
    A = 108838
    B = 111600
    node, developed_nodes_count, run_time, path_cost, Cf = RunCoolPathSearch(rm, A, B)
    import cPickle as pickle
    pickle.dump(node, open('cool.pkl','wb'))
